from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
from fastapi.responses import FileResponse
import io
import sys
import os
sys.path.append(os.path.dirname(__file__))

from model_npy import predict  # ← 本地使用相对路径导入


app = FastAPI()


app = FastAPI(
    title="Numpy Image Classifier API",
    description="上传 .npy 图像，返回二分类结果（CN 或 AD）",
    version="1.0.0"
)
from fastapi.middleware.cors import CORSMiddleware


# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)
@app.get("/.well-known/ai-plugin.json", include_in_schema=False)
def serve_ai_plugin():
    file_path = os.path.join(os.path.dirname(__file__), ".well-known", "ai-plugin.json")
    return FileResponse(path=file_path, media_type="application/json")



@app.post("/classify")
async def classify_npy(file: UploadFile = File(...)):
    """
    接收 .npy 文件作为输入图像，并返回预测结果。
    默认从 batch 中取第 0 张图像进行推理。
    """
    try:
        contents = await file.read()
        array = np.load(io.BytesIO(contents))

        if array.ndim not in [3, 4]:
            return JSONResponse(
                status_code=400,
                content={"error": f"输入维度不合法，应为 (H,W,C) 或 (B,H,W,C)，当前为 {array.shape}"}
            )

        result = predict(array)
        return JSONResponse(content=result)

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": f"处理失败: {str(e)}"}
        )

# 允许直接运行本地开发服务器
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
